﻿using System;
using System.Collections.Generic;
using CoreRay.Collections;

namespace CoreRay
{
    enum BvhCutPlane
    {
        PlaneX,
        PlaneY,
        PlaneZ
    }

    struct BvhCut
    {
        public BvhCut(BvhCutPlane cutPlane, float cutPosition) : this()
        {
            this.CutPlane = cutPlane;
            this.CutPosition = cutPosition;
        }

        public float CutPosition { get; set; }
        public BvhCutPlane CutPlane { get; set; }
    }

    public class BvhSpatial : SpatialStructure
    {
        private BinaryTree<BvhTreeNodeData> binaryTree;
        private List<GeometryObject> geometryObjects;
        private StatisticsItemValue maxDepthStat;
        private StatisticsItemMemory memoryStat;

        public BvhSpatial(bool useWorldMatrix, int minObjectsCount, int maxDepth) : base(useWorldMatrix)
        {
            this.MinObjectsCount = minObjectsCount;
            this.MaxDepth = maxDepth;
        }

        public int MinObjectsCount
        {
            get;
            private set;
        }

        public int MaxDepth 
        { 
            get; 
            private set; 
        }

        public override void Construct(List<GeometryObject> geometryObjects)
        {
            // Init statistics
            string statisticsCategory = "BVH Spatial - " + this.Name;

            this.maxDepthStat = new StatisticsItemValue("Max Depth");
            CoreRayContext.Current.Statistics.AddStatisticsItem(statisticsCategory, maxDepthStat);

            this.memoryStat = new StatisticsItemMemory("Memory Allocated");
            CoreRayContext.Current.Statistics.AddStatisticsItem(statisticsCategory, this.memoryStat);

            this.memoryStat.RecordMemoryUsageChange();

            // Init triangles list
            this.geometryObjects = geometryObjects;

            // Compute global bounding box
            this.BoundingBox = new BoundingBox();

            foreach(var geometryObject in this.geometryObjects)
            {
                this.BoundingBox.Add(this.UseWorldMatrix ? geometryObject.WorldBoundingBox : geometryObject.BoundingBox);
            }

            // Init binary tree
            this.binaryTree = new BinaryTree<BvhTreeNodeData>();
            this.binaryTree.Root = new BinaryTreeNode<BvhTreeNodeData>(new BvhTreeNodeData(this.BoundingBox));

            // Start recursive construction of the tree
            this.ConstructNode(this.binaryTree.Root, this.BoundingBox, 0, this.geometryObjects.Count, 0, 0);

            this.memoryStat.RecordMemoryUsageChange();
        }

        private void ConstructNode(BinaryTreeNode<BvhTreeNodeData> treeNode, BoundingBox boundingBox, int leftIndex, int rightIndex, int depth, int lastCutPlane)
        {
            // Update statistics
            if (this.maxDepthStat.Value < depth)
            {
                this.maxDepthStat.Value = depth;
            }

            int geometryObjectsCount = rightIndex - leftIndex;

            // Check to see if the current node has a minimum number of triangles
            if (geometryObjectsCount <= this.MinObjectsCount || depth >= this.MaxDepth)
            {
                treeNode.Data.IsLeafNode = true;
                treeNode.Data.TrianglesIndex = leftIndex;
                treeNode.Data.TrianglesCount = geometryObjectsCount;

                return;
            }

            // Find a cut point and a cut plane
            BvhCut cut = FindCutPlane(boundingBox, geometryObjectsCount, lastCutPlane);

            // Reorder the triangles list and separate them into a left part and a right part
            int leftCounter = 0;
            int rightCounter = geometryObjectsCount - 1;

            GeometryObject[] geometryObjectsBuffer = new GeometryObject[geometryObjectsCount];
            geometryObjects.CopyTo(leftIndex, geometryObjectsBuffer, 0, geometryObjectsCount);

            for (int i = 0; i < geometryObjectsCount; i++)
            {
                float boundingBoxAxisMaxPoint;

                if (cut.CutPlane == BvhCutPlane.PlaneX)
                {
                    boundingBoxAxisMaxPoint = this.UseWorldMatrix ? geometryObjectsBuffer[i].WorldBoundingBox.Center.X  : geometryObjectsBuffer[i].BoundingBox.Center.X;
                }

                else if (cut.CutPlane == BvhCutPlane.PlaneY)
                {
                    boundingBoxAxisMaxPoint = this.UseWorldMatrix ? geometryObjectsBuffer[i].WorldBoundingBox.Center.Y : geometryObjectsBuffer[i].BoundingBox.Center.Y;
                }

                else
                {
                    boundingBoxAxisMaxPoint = this.UseWorldMatrix ? geometryObjectsBuffer[i].WorldBoundingBox.Center.Z : geometryObjectsBuffer[i].BoundingBox.Center.Z;
                }

                // Triangle is on the left part of the cut
                if (boundingBoxAxisMaxPoint < cut.CutPosition)
                {
                    geometryObjects[leftIndex + leftCounter] = geometryObjectsBuffer[i];
                    leftCounter++;
                }

                else
                {
                    geometryObjects[leftIndex + rightCounter] = geometryObjectsBuffer[i];
                    rightCounter--;
                }
            }

            if (leftCounter == 0 || leftCounter == geometryObjectsCount)
            {
                /* Couldn't partition bboxes into left and right.
             * Force split at object median.
             */
                leftCounter = geometryObjectsCount / 2;
            }

            // Create left sub node
            BoundingBox leftBoundingBox = new BoundingBox();

            for (int i = leftIndex; i < leftIndex + leftCounter; i++)
            {
                leftBoundingBox.Add(this.UseWorldMatrix ? this.geometryObjects[i].WorldBoundingBox : this.geometryObjects[i].BoundingBox);
            }

            BvhTreeNodeData leftData = new BvhTreeNodeData(leftBoundingBox);
            treeNode.LeftChild = new BinaryTreeNode<BvhTreeNodeData>(leftData);

            this.ConstructNode(treeNode.LeftChild, treeNode.LeftChild.Data.BoundingBox, leftIndex, leftIndex + leftCounter, depth+1, (int) cut.CutPlane);

            // Create right sub node
            BoundingBox rightBoundingBox = new BoundingBox();

            for (int i = leftIndex + leftCounter; i < rightIndex; i++)
            {
                rightBoundingBox.Add(this.UseWorldMatrix ? this.geometryObjects[i].WorldBoundingBox : this.geometryObjects[i].BoundingBox);
            }

            BvhTreeNodeData rightData = new BvhTreeNodeData(rightBoundingBox);
            treeNode.RightChild = new BinaryTreeNode<BvhTreeNodeData>(rightData);

            this.ConstructNode(treeNode.RightChild, treeNode.RightChild.Data.BoundingBox, leftIndex + leftCounter, rightIndex, depth + 1, (int) cut.CutPlane);
        }

        public override Intersection Intersect(Ray ray)
        {
            Intersection intersection = IntersectNode(ray, this.binaryTree.Root);
            return intersection;
        }

        public override bool IntersectPosition(Ray ray)
        {
            return IntersectNodePosition(ray, this.binaryTree.Root);
        }

        private Intersection IntersectNode(Ray ray, BinaryTreeNode<BvhTreeNodeData> node)
        {
            Intersection intersection = null;
            
            if (node.Data.IsLeafNode)
            {
                for (int i = 0; i < node.Data.TrianglesCount; i++)
                {
                    if (this.UseWorldMatrix ? this.geometryObjects[node.Data.TrianglesIndex + i].WorldBoundingBox.Intersect(ray) : this.geometryObjects[node.Data.TrianglesIndex + i].BoundingBox.Intersect(ray))
                    {
                        Intersection result = this.geometryObjects[node.Data.TrianglesIndex + i].Intersect(ray);

                        if (result != null && (intersection == null || intersection.RayT > result.RayT))
                        {
                            intersection = result;
                        }
                    }
                }
            }

            else
            {
                float intersectionLeft = 0.0f;
                float intersectionRight = 0.0f;

                if (node.LeftChild != null)
                {
                    intersectionLeft = node.LeftChild.Data.BoundingBox.IntersectDistance(ray);
                }

                if (node.RightChild != null)
                {
                    intersectionRight = node.RightChild.Data.BoundingBox.IntersectDistance(ray);
                }

                if (intersectionLeft != 0.0f || intersectionRight != 0.0f)
                {
                    if (intersectionLeft != 0.0f && (intersectionLeft < intersectionRight || intersectionRight == 0.0f))
                    {
                        intersection = IntersectNode(ray, node.LeftChild);

                        if (intersectionRight != 0.0f && (intersection == null || intersection.RayT > intersectionRight))
                        {
                            Intersection intersectionNodeRight = IntersectNode(ray, node.RightChild);

                            if (intersectionNodeRight != null && (intersection == null || intersectionNodeRight.RayT < intersection.RayT))
                            {
                                intersection = intersectionNodeRight;
                            }
                        }
                    }

                    else if (intersectionRight != 0.0f)
                    {
                        intersection = IntersectNode(ray, node.RightChild);

                        if (intersectionLeft != 0.0f && (intersection == null || intersection.RayT > intersectionLeft))
                        {
                            Intersection intersectionNodeLeft = IntersectNode(ray, node.LeftChild);

                            if (intersectionNodeLeft != null && (intersection == null || intersectionNodeLeft.RayT < intersection.RayT))
                            {
                                intersection = intersectionNodeLeft;
                            }
                        }
                    }
                }
            }

            return intersection;
        }

        private bool IntersectNodePosition(Ray ray, BinaryTreeNode<BvhTreeNodeData> node)
        {
            bool intersection = false;

            if (node.Data.IsLeafNode)
            {
                for (int i = 0; i < node.Data.TrianglesCount; i++)
                {
                    if (this.UseWorldMatrix ? this.geometryObjects[node.Data.TrianglesIndex + i].WorldBoundingBox.Intersect(ray) : this.geometryObjects[node.Data.TrianglesIndex + i].BoundingBox.Intersect(ray))
                    {
                        bool result = this.geometryObjects[node.Data.TrianglesIndex + i].IntersectPosition(ray);

                        if (result == true)
                        {
                            return true;
                        }
                    }
                }
            }

            else
            {
                float intersectionLeft = 0.0f;
                float intersectionRight = 0.0f;

                if (node.LeftChild != null)
                {
                    intersectionLeft = node.LeftChild.Data.BoundingBox.IntersectDistance(ray);
                }

                if (node.RightChild != null)
                {
                    intersectionRight = node.RightChild.Data.BoundingBox.IntersectDistance(ray);
                }

                if (intersectionLeft != 0.0f || intersectionRight != 0.0f)
                {
                    if (intersectionLeft != 0.0f && (intersectionLeft < intersectionRight || intersectionRight == 0.0f))
                    {
                        intersection = IntersectNodePosition(ray, node.LeftChild);

                        if (intersection == false && intersectionRight != 0)
                        {
                            intersection = IntersectNodePosition(ray, node.RightChild);
                        }
                    }

                    else if (intersectionRight != 0.0f)
                    {
                        intersection = IntersectNodePosition(ray, node.RightChild);

                        if (intersection == false && intersectionLeft != 0)
                        {
                            intersection = IntersectNodePosition(ray, node.LeftChild);
                        }
                    }
                }
            }

            return intersection;
        }

        private BvhCut FindCutPlane(BoundingBox boundingBox, int trianglesCount, int currentCutPlane)
        {
            BvhCutPlane cutPlane = (BvhCutPlane) currentCutPlane;
            float cutPosition;

            if (cutPlane == BvhCutPlane.PlaneX)
            {
                cutPosition = boundingBox.MinPoint.X + (boundingBox.MaxPoint.X - boundingBox.MinPoint.X) * 0.5f;
            }

            else if (cutPlane == BvhCutPlane.PlaneY)
            {
                cutPosition = boundingBox.MinPoint.Y + (boundingBox.MaxPoint.Y - boundingBox.MinPoint.Y) * 0.5f;
            }

            else
            {
                cutPosition = boundingBox.MinPoint.Z + (boundingBox.MaxPoint.Z - boundingBox.MinPoint.Z) * 0.5f;
            }

            return new BvhCut(cutPlane, cutPosition);
        }
    }
}
